#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import re
import asyncio
import json
from typing import Dict, List, Any, Optional, Tuple
from openai import AsyncOpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter

logger = logging.getLogger(__name__)

class AIService:
    """AI服务
    
    处理AI相关的功能，如生成问答对、向量化文本等。
    """
    
    def __init__(self, config: Dict[str, Any]):
        """初始化AI服务
        
        Args:
            config: 配置字典
        """
        self.config = config
        ai_config = config["ai"]
        
        # 配置OpenAI API
        self.client = AsyncOpenAI(
            api_key=ai_config["api_key"],
            base_url=ai_config["api_base"]
        )
        
        self.model = ai_config["model"]
        self.temperature = ai_config["temperature"]
        self.max_tokens = ai_config["max_tokens"]
    
    async def generate_qa_by_text_split(
        self, 
        text: str, 
        system_name: str,
        chunk_size: int = 500,
        overlap: int = 50
    ) -> List[Dict[str, str]]:
        """通过文本分割生成问答对
        
        将文本分割为固定大小的块，不调用AI接口，直接将每个块转换为问答对。generate_qa_by_aigenerate_qa_by_ai
        其中问题为系统名称+块序号，答案为块内容。
        过滤掉目录内容和非实质性内容。
        
        Args:
            text: 文本内容
            system_name: 系统名称
            chunk_size: 文本块大小（字符数）
            overlap: 重叠大小（字符数）
            
        Returns:
            List[Dict[str, str]]: 问答对列表
        """
        logger.info(f"开始通过纯文本分割生成问答对，文本长度: {len(text)}, 块大小: {chunk_size}, 重叠大小: {overlap}")
        
        # 预处理文本，去除目录和引用链接
        text = self._preprocess_text(text)
        
        # 分割文本
        chunks = self._split_text(text, chunk_size, overlap)
        logger.info(f"文本分割完成，共 {len(chunks)} 个块")
        
        qa_pairs = []
        
        for idx, chunk in enumerate(chunks):
            try:
                # 过滤掉无意义的块
                if chunk.strip() and self._is_valid_chunk(chunk):
                    # 直接创建问答对，不调用AI
                    qa_pair = {
                        "Q": f"{system_name} 知识点 #{idx+1}",
                        "A": chunk,
                        "source": chunk
                    }
                    qa_pairs.append(qa_pair)
                
                if (idx + 1) % 50 == 0:
                    logger.info(f"已处理文本块 {idx + 1}/{len(chunks)}")
                
            except Exception as e:
                logger.error(f"处理文本块 {idx + 1} 时出错: {str(e)}")
            
        logger.info(f"文本分割生成问答对完成，共 {len(qa_pairs)} 个问答对")
        return qa_pairs
    
    async def generate_qa_by_ai(
        self, 
        text: str, 
        system_name: str,
        chunk_size: int = 2000,
        overlap: int = 100
    ) -> List[Dict[str, str]]:
        """通过AI生成问答对
        
        让AI根据文本内容生成问答对。支持大文本渐进式处理。
        
        Args:
            text: 文本内容
            system_name: 系统名称
            chunk_size: 文本块大小（字符数），默认2000
            overlap: 重叠大小（字符数），默认100
            
        Returns:
            List[Dict[str, str]]: 问答对列表
        """
        logger.info(f"开始通过AI生成问答对，文本长度: {len(text)}, 块大小: {chunk_size}, 重叠大小: {overlap}")
        
        # 预处理文本
        text = self._preprocess_text(text)
        
        # 使用文本分割器分割文本为多个小块
        chunks = self._split_text(text, chunk_size, overlap)
        
        # 为每个块生成可能的标题，用于提高上下文质量
        chunk_titles = await self._generate_chunk_titles(chunks)
        
        logger.info(f"文本分割完成，共 {len(chunks)} 个块")
            
        # 初始化结果集
        all_qa_pairs = []
        
        # 按批次处理文本块
        batch_size = 5
        for batch_idx in range(0, len(chunks), batch_size):
            batch_end = min(batch_idx + batch_size, len(chunks))
            batch_chunks = chunks[batch_idx:batch_end]
            batch_titles = chunk_titles[batch_idx:batch_end] if chunk_titles else None
            
            logger.info(f"处理批次 {batch_idx//batch_size + 1}/{(len(chunks)+batch_size-1)//batch_size}，" 
                       f"包含块 {batch_idx+1}-{batch_end}")
            
            # 并行处理批次内的每个块
            tasks = []
            for i, chunk in enumerate(batch_chunks):
                title = batch_titles[i] if batch_titles else None
                task = self._generate_qa_for_text_chunk(
                    chunk, 
                    system_name,
                    title=title,
                    chunk_index=batch_idx + i,
                    total_chunks=len(chunks)
                )
                tasks.append(task)
            
            # 等待批次内所有任务完成
            batch_results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 处理批次结果
            for i, result in enumerate(batch_results):
                if isinstance(result, Exception):
                    logger.error(f"处理文本块 {batch_idx + i + 1} 失败: {str(result)}")
                    continue
                    
                if result:
                    all_qa_pairs.extend(result)
                    logger.info(f"文本块 {batch_idx + i + 1} 生成了 {len(result)} 个问答对")
                
            # 简单的流量控制，避免API限制
            if batch_end < len(chunks):
                await asyncio.sleep(1)
        
        # 按照文本块顺序和问答对序号排序，确保与原文顺序一致
        all_qa_pairs.sort(key=lambda x: (x.get('chunk_index', 0), x.get('qa_index', 0)))
        
        # 移除临时排序字段
        for qa_pair in all_qa_pairs:
            if 'chunk_index' in qa_pair:
                del qa_pair['chunk_index']
            if 'qa_index' in qa_pair:
                del qa_pair['qa_index']
                
        logger.info(f"AI生成问答对完成，共 {len(all_qa_pairs)} 个问答对")
        return all_qa_pairs
    
    async def _generate_chunk_titles(self, chunks: List[str]) -> List[str]:
        """为每个文本块生成简短标题
        
        Args:
            chunks: 文本块列表
            
        Returns:
            List[str]: 标题列表
        """
        # 对于小型数据集可以跳过这一步
        if len(chunks) <= 5:
            return []
            
        titles = []
        for i, chunk in enumerate(chunks):
            # 简单方式：使用块的前50个字符作为标题
            preview = chunk[:50].strip().replace("\n", " ")
            if len(preview) > 0:
                titles.append(f"内容片段 {i+1}: {preview}...")
            else:
                titles.append(f"内容片段 {i+1}")
                
        return titles
        
    async def _generate_qa_for_text_chunk(
        self,
        chunk: str,
        system_name: str,
        title: Optional[str] = None,
        chunk_index: int = 0,
        total_chunks: int = 1,
        qa_per_chunk: int = 3  # 控制每个文本块生成的问答对数量
    ) -> List[Dict[str, str]]:
        """从文本块生成多个问答对
        
        Args:
            chunk: 文本块
            system_name: 系统名称
            title: 文本块标题（可选）
            chunk_index: 当前块索引
            total_chunks: 总块数
            qa_per_chunk: 每个块生成问答对的数量
            
        Returns:
            List[Dict[str, str]]: 问答对列表
        """
        if not chunk.strip() or not self._is_valid_chunk(chunk):
            logger.info(f"跳过无效文本块 {chunk_index+1}/{total_chunks}")
            return []
            
        context_info = ""
        if title:
            context_info = f"\n这是文档的一部分: {title}"
            
        # 更新提示词，要求输出JSON格式
        prompt = f"""根据以下{system_name}的文本内容，生成个高质量的问答对。
这些问答对将用于知识库构建，应该具有实用价值。

严格按照生成的问答对要求：
1. 提取核心操作步骤和功能说明作为问题
2. 答案需完整保留原始图片的Markdown格式链接
3. 自动过滤"如图*"的编号描述，不出现在答案内容中（如图4.1.1.1）
4. 不要使用制表符\\t或其他空白控制字符
5. 答案内容应该是连续的文本，避免不必要的空格和缩进
6. 严禁在答案末尾使用反斜杠\\作为续行符
7. 每个操作步骤生成独立的问答对
8. 每一个QA的结果是完整独立

文本内容:
{chunk}

请将结果以有效的JSON数组格式返回，格式如下：
```json
[
  {{
    "Q": "问题1",
    "A": "答案1（保留原始markdown格式）"
  }},
  {{
    "Q": "问题2",
    "A": "答案2（保留原始markdown格式）"
  }},
  ...
]
```

仅返回JSON数组，不要有其他说明文字。"""

        try:
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是专业的知识问答生成助手，擅长从文本中提取核心知识点并转化为高质量问答对。必须以JSON格式输出结果。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=max(0.1, self.temperature - 0.1),  # 略微降低温度以提高一致性
                max_tokens=min(2000, self.max_tokens),  # 限制token以避免超出限制
                response_format={"type": "json_object"}  # 强制返回JSON格式
            )
            
            # 解析JSON响应
            content = response.choices[0].message.content
            qa_pairs = self._parse_json_qa_response(content, chunk, chunk_index)
            
            # 记录生成的问答对
            if qa_pairs:
                logger.info(f"成功从文本块 {chunk_index+1}/{total_chunks} 生成 {len(qa_pairs)} 个问答对")
            else:
                logger.warning(f"文本块 {chunk_index+1}/{total_chunks} 未能生成有效问答对")
                
            return qa_pairs
            
        except Exception as e:
            logger.error(f"为文本块 {chunk_index+1}/{total_chunks} 生成问答对失败: {str(e)}")
            # 允许个别块失败，不中断整体流程
            return []
            
    def _parse_json_qa_response(
        self, 
        response: str,
        source_text: str,
        chunk_index: int = 0
    ) -> List[Dict[str, str]]:
        """解析AI返回的JSON格式问答对
        
        Args:
            response: AI响应内容（JSON格式）
            source_text: 源文本
            chunk_index: 块索引，用于标记问答对的顺序
            
        Returns:
            List[Dict[str, str]]: 问答对列表
        """
        qa_pairs = []
        
        try:
            # 尝试解析JSON
            # 可能需要清理一下响应内容，去除可能存在的markdown标记
            json_str = response
            if "```json" in response:
                json_str = re.search(r'```json\s*([\s\S]*?)\s*```', response).group(1)
            elif "```" in response:
                json_str = re.search(r'```\s*([\s\S]*?)\s*```', response).group(1)
                
            data = json.loads(json_str)
            
            # 根据返回数据结构处理
            if isinstance(data, list):
                items = data
            elif isinstance(data, dict) and "qa_pairs" in data:
                items = data["qa_pairs"]
            elif isinstance(data, dict) and "items" in data:
                items = data["items"]
            elif isinstance(data, dict):
                # 尝试找出包含问答对的字段
                for key, value in data.items():
                    if isinstance(value, list) and len(value) > 0:
                        items = value
                        break
                else:
                    items = [data]  # 可能整个对象就是一个问答对
            else:
                items = []
                
            # 处理每个问答对
            for i, item in enumerate(items):
                if not isinstance(item, dict):
                    continue
                    
                # 获取问题和答案字段
                Q = None
                A = None
                
                # 尝试不同的可能字段名
                for q_field in ["Q", "question", "问题", "q", "Question"]:
                    if q_field in item:
                        Q = item[q_field].strip()
                        break
                        
                for a_field in ["A", "answer", "答案", "a", "Answer"]:
                    if a_field in item:
                        A = item[a_field].strip()
                        break
                        
                if Q and A:
                    # 保留问号结尾
                    if not Q.endswith('?') and not Q.endswith('？'):
                        Q = Q + '?'
                        
                    # 去除答案末尾的反斜杠
                    A = A.rstrip('\\')
                        
                    # 构建问答对，包含块索引信息以保持顺序
                    qa_pairs.append({
                        "Q": Q,
                        "A": A,
                        "source": source_text[:100] + "..." if len(source_text) > 100 else source_text,
                        "chunk_index": chunk_index,
                        "qa_index": i
                    })
                    
            return qa_pairs
            
        except Exception as e:
            logger.error(f"解析JSON响应失败: {str(e)}")
            # 如果JSON解析失败，尝试使用正则表达式匹配
            return self._parse_qa_response_robust(response, source_text)

    def _parse_qa_response_robust(
        self, 
        response: str,
        source_text: str
    ) -> List[Dict[str, str]]:
        """增强版问答对响应解析
        
        使用多种模式匹配方法提取问答对，提高健壮性
        
        Args:
            response: AI响应内容
            source_text: 源文本
            
        Returns:
            List[Dict[str, str]]: 问答对列表
        """
        qa_pairs = []
        
        # 尝试多种模式匹配
        # 1. 使用编号+问题/答案匹配
        pattern1 = r'问题(\d+)[.:：]?\s*(.*?)\s*(?:\n|$).*?答案\1[.:：]?\s*(.*?)(?:\n问题\d+[.:：]|\n答案\d+[.:：]|$)'
        matches1 = re.finditer(pattern1, response, re.DOTALL)
        
        for match in matches1:
            Q = match.group(2).strip()
            A = match.group(3).strip()
            if Q and A:
                qa_pairs.append({
                    "Q": Q,
                    "A": A,
                    "source": source_text[:200] + "..." if len(source_text) > 200 else source_text
                })
        
        # 如果第一种模式没有匹配结果，尝试第二种模式
        if not qa_pairs:
            # 2. 简单问题/答案匹配
            pattern2 = r'问题[.:：]?\s*(.*?)\s*(?:\n|$).*?答案[.:：]?\s*(.*?)(?:\n问题[.:：]|\n答案[.:：]|$)'
            matches2 = re.finditer(pattern2, response, re.DOTALL)
            
            for match in matches2:
                Q = match.group(1).strip()
                A = match.group(2).strip()
                if Q and A:
                    qa_pairs.append({
                        "Q": Q,
                        "A": A,
                        "source": source_text[:200] + "..." if len(source_text) > 200 else source_text
                    })
        
        # 如果前两种都没有匹配结果，再尝试第三种方式
        if not qa_pairs:
            # 3. 按行分割，寻找问题和答案
            lines = response.split('\n')
            current_Q = None
            current_A = []
            
            for line in lines:
                q_match = re.match(r'^(.+\?)\s*$', line.strip())
                if q_match and not current_Q:
                    # 找到了问题行
                    current_Q = q_match.group(1)
                elif current_Q and line.strip() and not line.strip().endswith('?'):
                    # 累积答案行
                    current_A.append(line.strip())
                elif current_Q and current_A and (not line.strip() or q_match):
                    # 问答对结束
                    qa_pairs.append({
                        "Q": current_Q,
                        "A": " ".join(current_A),
                        "source": source_text[:200] + "..." if len(source_text) > 200 else source_text
                    })
                    
                    # 重置状态
                    current_Q = q_match.group(1) if q_match else None
                    current_A = []
            
            # 处理最后一个问答对
            if current_Q and current_A:
                qa_pairs.append({
                    "Q": current_Q,
                    "A": " ".join(current_A),
                    "source": source_text[:200] + "..." if len(source_text) > 200 else source_text
                })
        
        # 对结果进行简单清理
        for pair in qa_pairs:
            # 确保问题以问号结尾
            if not pair["Q"].endswith('?') and not pair["Q"].endswith('？'):
                pair["Q"] = pair["Q"] + '?'
            
            # 移除答案中可能的问题标记
            pair["A"] = re.sub(r'^答案[.:：]?\s*', '', pair["A"])
            
            # 去除答案末尾的反斜杠
            pair["A"] = pair["A"].rstrip('\\')
        
        return qa_pairs
    
    def _split_text(
        self, 
        text: str, 
        chunk_size: int = 500,
        overlap: int = 50
    ) -> List[str]:
        """将文本分割为固定大小的块
        
        使用langchain的RecursiveCharacterTextSplitter进行文本分割
        
        Args:
            text: 文本内容
            chunk_size: 块大小（字符数）
            overlap: 重叠大小（字符数）
            
        Returns:
            List[str]: 文本块列表
        """
        if not text:
            return []
        
        logger.info(f"使用langchain进行文本分割, 块大小: {chunk_size}, 重叠大小: {overlap}")
        
        # 使用langchain的RecursiveCharacterTextSplitter
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=overlap,
            length_function=len,
            separators=["\n\n", "\n", "。", "！", "？", ".", "!", "?", " ", ""]
        )
        
        # 执行分割
        chunks = text_splitter.split_text(text)
        
        logger.info(f"文本分割完成，共 {len(chunks)} 个块")
        return chunks
    
    def _parse_qa_response(
        self, 
        response: str,
        source_text: str
    ) -> List[Dict[str, str]]:
        """解析AI生成的问答对响应
        
        Args:
            response: AI的响应内容
            source_text: 源文本
            
        Returns:
            List[Dict[str, str]]: 问答对列表
        """
        qa_pairs = []
        
        # 分割成多个问答对
        pairs = re.split(r'\n\d+[\.、]', response)
        if len(pairs) == 1:  # 如果没有按编号分割成功，尝试按问题分割
            pairs = re.split(r'\n+问题[:：]', response)
            if len(pairs) > 1:
                # 修复第一个元素
                pairs = pairs[1:]
                pairs = [f"问题:{p}" for p in pairs]
        
        for pair in pairs:
            if not pair.strip():
                continue
                
            Q_match = re.search(r"问题[:：](.*?)(?:\n|$)", pair, re.DOTALL)
            A_match = re.search(r"答案[:：](.*?)(?:\n|$)", pair, re.DOTALL)
            
            if Q_match and A_match:
                Q = Q_match.group(1).strip()
                A = A_match.group(1).strip()
                # 去除答案末尾的反斜杠
                A = A.rstrip('\\')
                
                qa_pairs.append({
                    "Q": Q,
                    "A": A,
                    "source": source_text
                })
        
        return qa_pairs 
    
    def _preprocess_text(self, text: str) -> str:
        """预处理文本，去除目录和引用链接
        
        Args:
            text: 原始文本
            
        Returns:
            str: 处理后的文本
        """
        # 移除目录行（通常包含页码引用如 [xxx 123](#_Toc12345)）
        lines = text.split('\n')
        processed_lines = []
        
        # 标记是否在目录部分
        in_toc = False
        
        for line in lines:
            # 检测目录项特征
            is_toc_line = bool(re.search(r'\[\S.*\]\(#_Toc\d+\)', line))
            
            # 检测目录结束（通常是遇到章节标题如 # 引言）
            if in_toc and line.startswith('#'):
                in_toc = False
            
            # 如果发现目录行，标记为在目录部分
            if is_toc_line:
                in_toc = True
                continue
            
            # 不在目录部分的内容保留
            if not in_toc:
                processed_lines.append(line)
        
        return '\n'.join(processed_lines)
    
    def _is_valid_chunk(self, chunk: str) -> bool:
        """判断文本块是否有效（包含实质性内容）
        
        检查文本块是否包含非目录内容，以及是否有足够的实质性内容
        
        Args:
            chunk: 文本块
            
        Returns:
            bool: 是否为有效块
        """
        # 如果块太短，可能不包含有价值的信息
        if len(chunk.strip()) < 20:
            return False
        
        # 检查是否只包含目录项
        if re.search(r'\[\S.*\]\(#_Toc\d+\)', chunk) and len(chunk.split('\n')) < 3:
            return False
        
        # 检查是否只包含标题
        if chunk.strip().startswith('#') and len(chunk.split('\n')) < 2:
            return False
        
        # 检查是否包含足够的文本（不只是空格、标点符号等）
        text_content = re.sub(r'[^\w\s]', '', chunk)
        if len(text_content.strip()) < 15:
            return False
        
        return True 